import os
import torch
from PIL import Image
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from sklearn.model_selection import train_test_split


class MakeDataset():
    def __init__(self, imageDir, labelFile, transform=None):
        self.imageDir = imageDir
        self.labelFile = labelFile

        self.label = np.loadtxt(self.labelFile)
        self.transform = transform

    def __getitem__(self, idx):
        # 获取文件路径
        imgs_path = os.path.join(self.imageDir, "%.4d.jpg" % idx)
        image = Image.open(imgs_path)
        if self.transform is not None:
            image = self.transform(image)

        labels = self.label[idx]
        return image, labels

    def __len__(self):
        return self.label.shape[0]

    def make_dataset(self):
        train_rate = 0.7
        val_rate = 0.1
        test_rate = 0.2

        imgs, labels = self.getitem()

        X_train, X_test, y_train, y_test = train_test_split(imgs, labels, train_size=train_rate)
        X_val, X_test, y_val, y_test = train_test_split(X_test, y_test, train_size=val_rate / (val_rate + test_rate))


if __name__ == '__main__':
    DataDir = './image'
    dataset = MakeDataset(DataDir)
    dataset.getiten()
